

import sys
import datetime
import pickle
import cv2
import numpy as np

import glob
#coding : utf-8
import os
import abc
import xml.dom.minidom as xml


# # class_name = ["car", "van", "bus", "truck", "plate"]

np.set_printoptions(suppress=True)


class XmlReader(object):
    __metaclass__ = abc.ABCMeta

    def __init__(self):
        pass

    def read_content(self, filename):
        content = None
        if (False == os.path.exists(filename)):
            return content
        filehandle = None
        try:
            filehandle = open(filename, 'rb')
        except FileNotFoundError as e:
            print(e.strerror)
        try:
            content = filehandle.read()
        except IOError as e:
            print(e.strerror)
        if (None != filehandle):
            filehandle.close()
        if(None != content):
            return content.decode("utf-8", "ignore")
        return content

    @abc.abstractmethod
    def load(self, filename):
        pass


ROOT_PATH = "/home/data"  # os.getcwd()


img_s_sum = []
label_150 = []
img_cnt = [0]


class XmlTester(XmlReader):
    def __init__(self):
        XmlReader.__init__(self)

    def load(self, filename, img_id):
        filecontent = XmlReader.read_content(self, filename)
        img_map = dict()
        if None != filecontent:
            # print(filename)
            dom = xml.parseString(filecontent)
            im_size = dom.getElementsByTagName('size')[0]
            img = cv2.imread(filename[:-3]+"jpg")
            im_w = int((im_size.getElementsByTagName(
                'width')[0]).childNodes[0].data)
            im_h = int((im_size.getElementsByTagName(
                "height")[0]).childNodes[0].data)
            self.im_shape = np.array([im_w, im_h])
            self.bbox = []
            self.name = dom.getElementsByTagName(
                'filename')[0].childNodes[0].data
            self.obj_name = []
            for obj in dom.getElementsByTagName('object'):
                box = obj.getElementsByTagName('bndbox')[0]

                b_xmin = int((box.getElementsByTagName(
                    "xmin")[0]).childNodes[0].data)
                b_ymin = int((box.getElementsByTagName(
                    "ymin")[0]).childNodes[0].data)
                b_xmax = int((box.getElementsByTagName(
                    "xmax")[0]).childNodes[0].data)
                b_ymax = int((box.getElementsByTagName(
                    "ymax")[0]).childNodes[0].data)
                b_name = obj.getElementsByTagName(
                    "name")[0].childNodes[0].data

                self.obj_name.append(b_name)
                self.bbox.append([[b_xmin, b_ymin], [b_xmax, b_ymax]])
                if b_name[-2:] == "open_bed_heavy_truck":
                    os.system(f"mkdir -p data/train/{b_name}")
                    attr1 = obj.getElementsByTagName("attributes")[0]
                    # attr2 =attr1.getElementsByTagName("attribute")[1] #[0].childNodes[0].data
                    car_color = attr1.getElementsByTagName(
                        "value")[1].childNodes[0].data
                    if car_color.isdigit():
                        id_ = car_color
                        car_color = attr1.getElementsByTagName(
                            "value")[0].childNodes[0].data
                    else:
                        id_ = attr1.getElementsByTagName(
                            "value")[0].childNodes[0].data
                    img0 = img[(b_ymin):(b_ymax),
                               (b_xmin):(b_xmax)]
                    # print(img0.shape, mask0.shape)
                    if (img is None) or img0 is None or len(img0.shape) != 3:
                        print("空图片1", img0 is None)
                        continue
                    if (b_ymin >= (b_ymax)) or ((b_xmin) >= (b_xmax)):
                        print("空图片2", self.bbox[-1])
                        continue
                    if (mask0.shape[0] < 100) or (mask0.shape[1] < 100):
                        print("图片小于100像素")
                        continue
                    path = f"./data/train/{b_name}/{img_cnt[0]}.bmp"
                    cv2.imwrite(path, img0)
                    img_cnt[0] += 1
                elif b_name[-2:] == "ck":
                    os.system(f"mkdir -p data/train/{b_name}")

                    img0 = img[(b_ymin):(b_ymax),
                               (b_xmin):(b_xmax)]
                    # print(img0.shape, mask0.shape)
                    if (img is None) or img0 is None or len(img0.shape) != 3:
                        print("空图片1", img0 is None)
                        continue
                    if (b_ymin >= (b_ymax)) or ((b_xmin) >= (b_xmax)):
                        print("空图片2", self.bbox[-1])
                        continue
                    path = f"./data/train/{b_name}/{img_cnt[0]}.bmp"
                    cv2.imwrite(path, img0)
                    img_cnt[0] += 1
            self.polygon = []
            return self.im_shape, None

    def xyxy2yolo(self):
        self.bboxyolo = []
        for bx in self.bbox:  # xyxy
            cent = (bx[1] + bx[0]) / 2
            dwh = bx[1] - bx[0]
            # self.bboxyolo.append([bx[0]/self.im_shape,dwh/self.im_shape])
            self.bboxyolo.append(
                [cent / self.im_shape, dwh / self.im_shape])  # cxcywh
        return np.array(self.bboxyolo).reshape(-1, 4), self.obj_name


def to_color_2d():

    xmls = glob.glob("%s/1441/*.xml" % ROOT_PATH)
    for i, xml_path in enumerate(xmls):
        reader = XmlTester()
        im_shape, bbox = reader.load(xml_path, True)
    # wb是覆盖写，如果需要追加，则为‘ab'
    datas = {'img': img_s_sum, 'label': label_150}
    f = open('data.pkl', 'wb')
    data = pickle.dump(datas, f, -1)
    f.close()
    print("img len:", len(img_s_sum))


def had(fn, class_name):
    xmls = glob.glob(f"{ROOT_PATH}/{fn}/*.xml")
    cont = class_name.copy()
    for k in cont:
        cont[k] = 0
    for i, xml_path in enumerate(xmls):
        reader = XmlTester()
        _, img_map = reader.load(xml_path, i)


def to_yolo_2d_fire():
    class_name = {'fire': 0, 'big_fire': 1, 'smoke': 2, 'others': 3}
    had("599", class_name)
    had("1225", class_name)
    had("820", class_name)


def to_yolo_2d_zhan():
    class_name = {'fixed_stall': 0, 'sunshade': 1}
    had("1246", class_name)
    had("651", class_name)


def to_yolo_2d_dianti():
    class_name = {'electric_scooter': 0,
                  'person': 1, 'bicycle': 2, 'others': 3}

    had("1227", class_name)
    had("1395", class_name)


def to_yolo_2d_car():
    class_name = {'car': 0}
    os.system("mkdir -p /home/data/700")
    xmls = glob.glob(f"{ROOT_PATH}/1583/*.xml")
    for i, xml_path in enumerate(xmls):
        if i % 3 == 1:
            os.system(f"cp {xml_path} {ROOT_PATH}/700")
            os.system(f"cp {(xml_path[:-3]+'jpg')} {ROOT_PATH}/700")
    had("1583", class_name)
    had("700", class_name)


def to_yolo_2d_person():
    class_name = {'person': 0}
    os.system("mkdir -p /home/data/700")
    xmls = glob.glob(f"{ROOT_PATH}/757/*.xml")
    for i, xml_path in enumerate(xmls):
        if i % 10 == 1:
            os.system(f"mv {xml_path} {ROOT_PATH}/700")
            os.system(f"mv {(xml_path[:-3]+'jpg')} {ROOT_PATH}/700")
    had("757", class_name)
    had("700", class_name)


def to_yolo_2d_zhatu():
    # class_name = {"car":0, "van":1, "bus":2, "open_bed_heavy_truck":3, "close_bed_heavy_truck":4,
    #               "open_bed_light_truck":5, "close_bed_light_truck":6, "others_truck":7}
    class_name = {"dirt": 0, "truck_bed_cover": 1,
                  "empty_bed": 2, "side_plate": 3, "back_plate": 4}
    had("1441", class_name)
    # had("1442", class_name)


def to_yolo_2d_chushi():
    class_name = {"person": 30, "short_sleeve_red": 0, "short_sleeve_black": 1, "short_sleeve_white": 2, "short_sleeve_grey": 3, "short_sleeve_green": 4                  }
    had("1593", class_name)


funs = {"zhan": to_yolo_2d_zhan, "car": to_yolo_2d_car, "fire": to_yolo_2d_fire, "dianti": to_yolo_2d_dianti,
        "person": to_yolo_2d_person, "zhatu": to_yolo_2d_zhatu, "chushi": to_yolo_2d_chushi}
start_time = datetime.datetime.now()  # 记录程序开始执行的当前时间
print("脚本名", sys.argv[0], str(sys.argv[1]))

funs[str(sys.argv[1])]()

stop_time = datetime.datetime.now()  # 记录执行结束的当前时间
func_time = stop_time - start_time  # 得到中间功能的运行时间
print("func is running %s s" % func_time)
